import base64
import logging
import smtplib
from email.mime.text import MIMEText
from typing import TypedDict, Any

import toml
import uvicorn
from fastapi import FastAPI
from langchain_core.messages import HumanMessage, AIMessage
from langchain_core.tools import tool
from langchain_openai import ChatOpenAI
from langgraph.checkpoint.memory import MemorySaver
from langgraph.prebuilt import create_react_agent
from pydantic import BaseModel

# --- Logging Setup ---
logging.basicConfig(
    level=logging.INFO,
    format='%(asctime)s - %(name)s - %(levelname)s - %(message)s'
)
logger = logging.getLogger(__name__)

# --- Configuration Loading ---
try:
    config = toml.load("config.toml")
    openai_config = config['openai']
    api_config = config['api']
    mcp_email_config = config['mcp_email']
    memory_config = config['memory']
    logger.info("Configuration loaded successfully.")
except (FileNotFoundError, KeyError) as e:
    logger.error(f"Error: Configuration file 'config.toml' is missing or malformed. {e}")
    exit()


# --- 1. Tool Definition ---
class OnboardingDetails(TypedDict):
    """The required details for onboarding a new micro-enterprise client."""
    company_name: str
    abn: str
    contact_phone: str
    contact_email: str


@tool
def send_onboarding_invitation(details: OnboardingDetails) -> str:
    """
    Sends an onboarding invitation email to a new micro-enterprise client
    once all required information has been collected.
    """
    try:
        sender_email = mcp_email_config['user']
        receiver_email = details['contact_email']
        password = mcp_email_config.get("password")

        if not password:
            raise ValueError("Email password not found in config.toml.")

        subject = f"Invitation to Onboard: {details['company_name']}"
        body = f"""
        Hello,

        This is an invitation to onboard your company, {details['company_name']} (ABN: {details['abn']}), to our platform.

        Please click the following link to begin the process:
        https://onboard.example.com/start?session_id=UNIQUE_ID_HERE

        If you have any questions, please contact us.

        Best regards,
        Your Accounting Firm
        """

        msg = MIMEText(body)
        msg['Subject'] = subject
        msg['From'] = sender_email
        msg['To'] = receiver_email

        logger.info(f"Attempting to send invitation email to {receiver_email}")
        with smtplib.SMTP(mcp_email_config['host'], mcp_email_config['port']) as server:
            server.starttls()
            server.login(sender_email, password)
            server.send_message(msg)

        logger.info(f"Successfully sent invitation email to {receiver_email}.")
        return f"Successfully sent invitation email to {receiver_email}."

    except Exception as e:
        logger.error(f"Failed to send email: {e}", exc_info=True)
        return f"Failed to send email. Error: {e}"


# --- 2. Agent Setup ---
system_prompt = """你是一个税务代理的得力助手。
你的主要任务是为新的微小企业客户启动Onboarding（入驻）流程。

**对话流程与规则:**

1. **自我介绍与告知:**
   - 对话开始时，你必须先进行自我介绍，并清楚地告知用户你的功能和需要收集的全部信息。
   - 请这样说：“您好！我是您的智能入驻助手，负责为您的新客户启动入驻流程。我需要您提供以下四项信息：1. 公司名称, 2. ABN（澳大利亚商业编号）, 3. 联系人电话号码, 4. 联系人电子邮箱地址。”

2. **收集信息:**
   - 在介绍完毕后，再开始逐一收集信息。可以从“请问公司名称是什么？”开始。

3. **汇总与请求确认:**
   - 当你收集齐所有四项信息后，你必须将收集到的信息汇总展示给用户，然后明确地请求确认。
   - 请这样说：“好的，我已经收集到以下信息：\n- 公司名称: [公司名称]\n- ABN: [ABN]\n- 联系电话: [电话]\n- 联系邮箱: [邮箱]\n请问信息是否正确，可以发送邀请邮件了吗？”
   - 在这个阶段，你必须等待用户的回复，【绝对不能】调用工具。

4. **执行动作 - 关键规则与行为范例:**
   - 这是最重要的规则。在你发出确认请求后，如果用户的下一条消息是肯定的（例如“是的”、“确认”、“可以”），你的【唯一输出】必须是工具调用，【绝对不能】包含任何文字回复。
   - **请严格遵循以下行为范例:**
     ```
     AI: 好的，我已经收集到以下信息：
     - 公司名称: ABC Pty Ltd
     - ABN: 123456789
     - 联系电话: 0400123456
     - 联系邮箱: contact@abc.com
     请问信息是否正确，可以发送邀请邮件了吗？

     用户: 是的，信息正确，请发送。

     AI: [此处应直接调用 send_onboarding_invitation 工具，而不是生成文字]
     ```
   - 你的思考过程应该是：“用户已确认，现在必须调用工具。” 然后直接输出工具调用指令。

5. **结束对话:**
   - 成功调用工具后，你的最终回复应该是类似于“好的，邀请邮件已成功发送。感谢您的使用，再见！”来结束对话。
"""

model = ChatOpenAI(
    base_url=openai_config.get("base_url"),
    model=openai_config.get("model"),
    api_key=openai_config.get("api_key")
)
tools = [send_onboarding_invitation]


# --- Custom Checkpointer for Sliding Window Memory ---
class SlidingWindowMemory(MemorySaver):
    def __init__(self, *, k: int):
        super().__init__()
        self.k = k

    def put(self, m_config: dict, checkpoint: dict, metadata: dict, new_versions: dict) -> None:
        if 'messages' in checkpoint and len(checkpoint['messages']) > self.k:
            thread_id = m_config.get("configurable", {}).get("thread_id")
            logger.info(
                f"Trimming memory for thread '{thread_id}' before saving. From {len(checkpoint['messages'])} to {self.k} messages."
            )
            checkpoint['messages'] = checkpoint['messages'][-self.k:]
        super().put(m_config, checkpoint, metadata, new_versions)


window_size = memory_config["window_size"]
memory = SlidingWindowMemory(k=window_size)
logger.info(f"Memory configured with a sliding window of size {window_size}.")

# --- 3. Agent ---
app = create_react_agent(
    name="onboarding agent",
    model=model,
    tools=tools,
    checkpointer=memory,
    prompt=system_prompt
)

# --- 4. FastAPI Application ---
api_app = FastAPI(
    title="Onboarding Agent API",
    description="An API for interacting with the micro-enterprise onboarding assistant.",
    version="1.0.0"
)


class ChatRequest(BaseModel):
    conversation_id: str
    content: str


class ChatResponse(BaseModel):
    code: int
    error_msg: str
    data: Any


@api_app.post("/chat", response_model=ChatResponse)
def chat_with_agent(request: ChatRequest):
    """
    Handles a single turn in a conversation with the onboarding agent.
    """
    conv_id = request.conversation_id

    logger.info(f"Received request for conversation_id: '{conv_id}', content: '{request.content}'")

    try:
        decoded_content = base64.b64decode(request.content).decode("utf-8")
    except Exception as e:
        logging.error(f"Failed to decode base64 content: {e}")
        return ChatResponse(code=-1, error_msg="Failed to decode base64 content.", data={})

    if conv_id is None:
        return ChatResponse(code=-1, error_msg="No conversation ID provided.")

    inputs = {"messages": [HumanMessage(content=decoded_content)]}

    final_state = app.invoke(inputs, config={"configurable": {"thread_id": conv_id}})

    last_ai_message = next(
        (msg for msg in reversed(final_state['messages']) if isinstance(msg, AIMessage)),
        None
    )

    final_answer = last_ai_message.content if last_ai_message else "Sorry, I encountered an error."

    response = {
        "final_answer": final_answer,
        "conversation_id": conv_id,
    }

    logger.info(f"Sending response for conversation_id: '{conv_id}', response: '{response}'")
    return ChatResponse(code=0, error_msg="", data=response)


@api_app.get("/")
def read_root():
    return {"message": "Onboarding Agent API is running. Go to /docs for API documentation."}


# --- 5. Main Execution ---
if __name__ == "__main__":
    uvicorn.run(api_app, host=api_config["host"], port=api_config["port"])
